package complexdatastructure.multidatastructure;

import java.util.LinkedList;
import java.util.TreeMap;

/**
 * @Classname : MaxStack
 * @Description : <a href="https://leetcode.cn/problems/max-stack/">716. 最大栈</a>
 * @Author : chentianyu
 * @Date 2023/2/15 21:46
 */


public class MaxStack {
    // 双向链表节点
    class Node {
        int val;
        Node pre;
        Node next;
        public Node() {}
        public Node(int val) {this.val = val;}
    }

    // 双向链表
    class DLinkedList {
        Node head;
        Node tail;
        public DLinkedList() {
            head = new Node();
            tail = new Node();
            head.next = tail;
            tail.pre = head;
        }

        public Node push(int val) {
            Node node = new Node(val);
            Node pre = tail.pre;
            pre.next = node;
            node.pre = pre;
            node.next = tail;
            tail.pre = node;
            return node;
        }

        public int peek() {
            return tail.pre.val;
        }

        public int pop() {
            return unLink(tail.pre).val;
        }

        public Node unLink(Node node) {
            node.pre.next = node.next;
            node.next.pre = node.pre;
            node.pre = null;
            node.next = null;
            return node;
        }
    }

    private DLinkedList stack;  // 双向链表记录元素push顺序
    private TreeMap<Integer, LinkedList<Node>> map;  // 有序map记录已push元素节点的索引

    public MaxStack() {
        stack = new DLinkedList();
        map = new TreeMap<>();
    }

    public void push(int x) {
        Node node = stack.push(x);
        map.computeIfAbsent(x, key->new LinkedList<>()).addLast(node);
    }

    public int pop() {
        int key = stack.pop();
        map.get(key).removeLast();
        if (map.get(key).isEmpty()) map.remove(key);
        return key;
    }

    public int top() {
        return stack.peek();
    }

    public int peekMax() {
        return map.lastKey();
    }

    public int popMax() {
        int key = peekMax();
        Node node = map.get(key).removeLast();
        if (map.get(key).isEmpty()) map.remove(key);
        stack.unLink(node);
        return key;
    }
}
